package com.example.chat.agent.service.impl;

import com.example.chat.agent.model.ZhipuModel;
import com.example.chat.agent.service.AgentService;
import com.example.chat.agent.service.ChromaDbService;
import com.example.chat.common.constant.ChromaStoreContant;
import com.example.chat.common.constant.DigitConstant;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.List;

@Service
@Slf4j
public class AgentServiceImpl implements AgentService {
    @Resource
    ChromaDbService chromaDbService;
    @Resource
    ZhipuModel zhipuModel;

    @Override
    public String search(String question) {
        ChromaEmbeddingStore embeddingStore = chromaDbService.getEmbeddingStore(ChromaStoreContant.ZHIPU_TEXT);
        ZhipuAiEmbeddingModel embeddingModel = chromaDbService.getEmbeddingModel();

        Embedding queryEmbedding = embeddingModel.embed(question).content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder().queryEmbedding(queryEmbedding).maxResults(DigitConstant.THREE).build();
        EmbeddingSearchResult<TextSegment> embeddedEmbeddingSearchResult = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> embeddingMatcheList = embeddedEmbeddingSearchResult.matches();
        EmbeddingMatch<TextSegment> embeddingMatch = embeddingMatcheList.get(0);
        String text = embeddingMatch.embedded().text();

        String info = "问题："+question+"\\n下面给出一段内容，从内容中找出关于我提的问题的答案(注意你是一个机器人，去除一些不相干的信息，不用重复问题，只给我答案就行)，如果内容中未提及，就用大模型回答\\n" + text;
        UserMessage userMessage = UserMessage.userMessage(info);
        Response<AiMessage> response = zhipuModel.getZhipuAiChatModel().generate(userMessage);
        AiMessage aiMessage = response.content();
        log.info(aiMessage.text());

        return aiMessage.text();
    }
}
